#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Full PBPK Model for HSK7653

Comprehensive physiologically-based pharmacokinetic model for HSK7653
based on Study HSK7653-101-5 mg data.

Author: PBPK Modeling System
Version: 1.0.0
Date: 2025-07-23
"""

import json
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, Any, List
import os
import warnings
warnings.filterwarnings('ignore')

class FullPBPKHSK7653:
    """Full PBPK model for HSK7653 compound"""
    
    def __init__(self, parameter_file: str = None):
        """Initialize with HSK7653 parameters"""
        if parameter_file is None:
            # Use absolute path to find the parameter file in sample folder
            current_dir = os.path.dirname(os.path.abspath(__file__))
            parent_dir = os.path.dirname(current_dir)
            parameter_file = os.path.join(parent_dir, "data", "sample", "hsk7653_pbpk_parameters.json")
        
        self.parameter_file = parameter_file
        self.params = self._load_parameters()
        self.compartments = self._setup_compartments()
        self.time_points = None
        self.results = {}
    
    def _load_parameters(self) -> Dict[str, Any]:
        """Load HSK7653 parameters from JSON"""
        try:
            with open(self.parameter_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except FileNotFoundError:
            print(f"HSK7653 parameter file not found: {self.parameter_file}")
            raise
    
    def _setup_compartments(self) -> Dict[str, Dict[str, float]]:
        """Setup PBPK compartments with physiological parameters"""
        
        # Get body weight for scaling
        body_weight = self.params['demographics']['weight']
        
        # Scale organ volumes to actual body weight
        compartments = {
            'gut_lumen': {
                'volume': 0.5,  # L
                'flow': 0,      # L/h
                'kp': 1.0
            },
            'stomach': {
                'volume': 0.1,
                'flow': 0.25,
                'kp': 1.5
            },
            'duodenum': {
                'volume': 0.1,
                'flow': 0.35,
                'kp': 1.8
            },
            'jejunum': {
                'volume': 0.2,
                'flow': 0.7,
                'kp': 2.0
            },
            'ileum': {
                'volume': 0.3,
                'flow': 0.7,
                'kp': 2.2
            },
            'colon': {
                'volume': 0.5,
                'flow': 0.2,
                'kp': 1.0
            },
            'liver': {
                'volume': self.params['organ_distribution']['liver']['volume'],
                'flow': self.params['blood_flows']['hepatic_total'],
                'kp': self.params['organ_distribution']['liver']['kp']
            },
            'kidney': {
                'volume': self.params['organ_distribution']['kidney']['volume'],
                'flow': self.params['blood_flows']['renal'],
                'kp': self.params['organ_distribution']['kidney']['kp']
            },
            'brain': {
                'volume': self.params['organ_distribution']['brain']['volume'],
                'flow': self.params['blood_flows']['cerebral'],
                'kp': self.params['organ_distribution']['brain']['kp']
            },
            'lung': {
                'volume': self.params['organ_distribution']['lung']['volume'],
                'flow': self.params['blood_flows']['hepatic_arterial'],
                'kp': self.params['organ_distribution']['lung']['kp']
            },
            'heart': {
                'volume': self.params['organ_distribution']['heart']['volume'],
                'flow': self.params['blood_flows']['cardiac_output'] * 0.046,
                'kp': self.params['organ_distribution']['heart']['kp']
            },
            'muscle': {
                'volume': self.params['organ_distribution']['muscle']['volume'],
                'flow': self.params['blood_flows']['muscle'],
                'kp': self.params['organ_distribution']['muscle']['kp']
            },
            'adipose': {
                'volume': self.params['organ_distribution']['adipose']['volume'],
                'flow': self.params['blood_flows']['adipose'],
                'kp': self.params['organ_distribution']['adipose']['kp']
            },
            'spleen': {
                'volume': self.params['organ_distribution']['spleen']['volume'],
                'flow': self.params['blood_flows']['cardiac_output'] * 0.023,
                'kp': self.params['organ_distribution']['spleen']['kp']
            },
            'plasma': {
                'volume': self.params['reference_values']['plasma_volume'],
                'flow': self.params['blood_flows']['cardiac_output'],
                'kp': 1.0
            }
        }
        
        return compartments
    
    def calculate_absorption(self, dose: float, time_span: float = 24.0) -> Dict[str, Any]:
        """Calculate absorption using ADAM-like model"""
        
        time_points = np.linspace(0, time_span, 100)
        dt = time_points[1] - time_points[0]
        
        # Absorption parameters
        fa = self.params['pharmacokinetics']['absorption']['fa']
        fg = self.params['pharmacokinetics']['absorption']['fg']
        fh = self.params['pharmacokinetics']['absorption']['fh']
        ka = self.params['pharmacokinetics']['absorption']['ka']
        
        # Scale absorption based on realistic parameters
        # Use fraction absorbed instead of absolute amount
        fraction_absorbed = np.zeros(len(time_points))
        absorption_rate = np.zeros(len(time_points))
        
        # Maximum fraction that can be absorbed
        max_fraction = fa * fg * fh
        
        for i, t in enumerate(time_points):
            fraction_absorbed[i] = max_fraction * (1 - np.exp(-ka * t))
            if i > 0:
                absorption_rate[i] = (fraction_absorbed[i] - fraction_absorbed[i-1]) / dt * dose
            else:
                absorption_rate[i] = max_fraction * ka * dose
        
        return {
            'time_points': time_points,
            'absorbed_amount': fraction_absorbed * dose,
            'absorption_rate': absorption_rate,
            'total_absorbed': fraction_absorbed[-1] * dose
        }
    
    def run_full_pbpk_simulation(self, dose: float = None, time_span: float = 48.0) -> Dict[str, Any]:
        """Run complete Full PBPK simulation with proper compartmental modeling"""
        
        if dose is None:
            dose = self.params['drug_info']['dose']
        
        # Setup
        time_points = np.linspace(0, time_span, 200)
        dt = time_points[1] - time_points[0]
        
        # Get absorption data
        absorption_data = self.calculate_absorption(dose, time_span)
        
        # Initialize compartments
        concentrations = {}
        amounts = {}
        
        for comp_name, comp_data in self.compartments.items():
            concentrations[comp_name] = np.zeros(len(time_points))
            amounts[comp_name] = np.zeros(len(time_points))
        
        # Initial conditions - dose in mg
        gut_volume = self.compartments['gut_lumen']['volume']
        concentrations['gut_lumen'][0] = dose / gut_volume  # mg/L
        
        # Get key parameters
        plasma_volume = self.compartments['plasma']['volume']
        cardiac_output = self.params['blood_flows']['cardiac_output']
        
        # Simulation with proper mass balance
        for i in range(1, len(time_points)):
            # Absorption from gut to plasma
            if i < len(absorption_data['absorption_rate']):
                absorption_rate = absorption_data['absorption_rate'][i]  # mg/h
            else:
                absorption_rate = 0
            
            # Update plasma (central compartment)
            plasma_conc = concentrations['plasma'][i-1]
            
            # Input from absorption
            input_from_absorption = absorption_rate / plasma_volume
            
            # Elimination (systemic clearance)
            elimination = (self.params['pharmacokinetics']['clearance']['systemic'] * 
                          plasma_conc / plasma_volume)
            
            # Distribution to organs
            distribution_loss = 0
            for comp_name, comp_data in self.compartments.items():
                if comp_name not in ['gut_lumen', 'plasma']:
                    flow = comp_data['flow']
                    kp = comp_data['kp']
                    organ_conc = concentrations[comp_name][i-1]
                    
                    # Mass transfer: flow * (C_plasma/kp - C_organ) / organ_volume
                    transfer_rate = (flow * (plasma_conc/kp - organ_conc)) / plasma_volume
                    distribution_loss += transfer_rate
            
            # Update plasma concentration
            d_plasma = input_from_absorption - elimination - distribution_loss
            concentrations['plasma'][i] = max(0, plasma_conc + d_plasma * dt)
            
            # Update other organs
            for comp_name, comp_data in self.compartments.items():
                if comp_name not in ['gut_lumen', 'plasma']:
                    flow = comp_data['flow']
                    kp = comp_data['kp']
                    organ_volume = comp_data['volume']
                    organ_conc = concentrations[comp_name][i-1]
                    
                    # Mass transfer
                    transfer_rate = (flow * (plasma_conc/kp - organ_conc)) / organ_volume
                    concentrations[comp_name][i] = max(0, organ_conc + transfer_rate * dt)
            
            # Update gut lumen (depletion only)
            if i < len(absorption_data['absorption_rate']):
                gut_depletion = absorption_data['absorption_rate'][i] / gut_volume
                concentrations['gut_lumen'][i] = max(0, concentrations['gut_lumen'][i-1] - gut_depletion * dt)
            else:
                concentrations['gut_lumen'][i] = 0
        
        # Calculate derived parameters - concentrations already in mg/L
        plasma_auc = np.trapezoid(concentrations['plasma'], time_points)
        cmax = np.max(concentrations['plasma']) if len(concentrations['plasma']) > 0 else 0
        tmax_idx = np.argmax(concentrations['plasma']) if len(concentrations['plasma']) > 0 else 0
        tmax = time_points[tmax_idx] if len(time_points) > tmax_idx else 0
        
        # Calculate organ exposures
        organ_exposure = {}
        for comp_name in self.compartments.keys():
            if comp_name != 'gut_lumen':
                organ_exposure[comp_name] = np.trapezoid(concentrations[comp_name], time_points)
        
        return {
            'time_points': time_points,
            'concentrations': concentrations,
            'amounts': amounts,
            'plasma_auc': plasma_auc,
            'cmax': cmax,
            'tmax': tmax,
            'organ_exposure': organ_exposure,
            'absorption_data': absorption_data
        }
    
    def calculate_key_metrics(self, results: Dict[str, Any]) -> Dict[str, Any]:
        """Calculate key PK metrics"""
        
        plasma_conc = results['concentrations']['plasma']
        time_points = results['time_points']
        
        # AUC calculations
        auc_0_t = np.trapezoid(plasma_conc, time_points)
        
        # Terminal half-life
        terminal_points = plasma_conc[-10:]  # Last 10 points
        if np.any(terminal_points > 0):
            log_conc = np.log(terminal_points[terminal_points > 0])
            time_term = time_points[-len(terminal_points):][terminal_points > 0]
            if len(log_conc) > 1:
                slope, _ = np.polyfit(time_term, log_conc, 1)
                terminal_half_life = -0.693 / slope
            else:
                terminal_half_life = results['tmax'] * 2
        else:
            terminal_half_life = results['tmax'] * 2
        
        return {
            'auc': auc_0_t,
            'cmax': results['cmax'],
            'tmax': results['tmax'],
            'terminal_half_life': terminal_half_life,
            'clearance': self.params['drug_info']['dose'] / auc_0_t,
            'vd': (self.params['drug_info']['dose'] / results['cmax']) if results['cmax'] > 0 else 0
        }
    
    def generate_comprehensive_plots(self, results: Dict[str, Any], save_path: str = "../output"):
        """Generate comprehensive visualization for HSK7653"""
        
        os.makedirs(save_path, exist_ok=True)
        
        # Create figure with subplots
        fig = plt.figure(figsize=(24, 18))
        
        # Color scheme
        colors = plt.cm.Set2(np.linspace(0, 1, 10))
        time_points = results['time_points']
        
        # 1. Plasma concentration-time profile
        ax1 = plt.subplot(4, 3, 1)
        plasma_conc = results['concentrations']['plasma']
        ax1.plot(time_points, plasma_conc, color=colors[0], linewidth=3, label='HSK7653')
        ax1.set_xlabel('Time (hours)')
        ax1.set_ylabel('Plasma Concentration (mg/L)')
        ax1.set_title('HSK7653 Plasma Concentration-Time Profile')
        ax1.grid(True, alpha=0.3)
        ax1.legend()
        
        # 2. Log-scale plasma concentration
        ax2 = plt.subplot(4, 3, 2)
        mask = plasma_conc > 0
        if np.any(mask):
            ax2.semilogy(time_points[mask], plasma_conc[mask], color=colors[0], linewidth=3)
            ax2.set_xlabel('Time (hours)')
            ax2.set_ylabel('Plasma Concentration (mg/L) - Log Scale')
            ax2.set_title('Log-Linear Plasma Concentration')
            ax2.grid(True, alpha=0.3)
        
        # 3. Organ concentration comparison
        ax3 = plt.subplot(4, 3, 3)
        organs = ['liver', 'kidney', 'brain', 'lung', 'muscle', 'adipose', 'heart', 'spleen']
        final_conc = [results['concentrations'][org][-1] for org in organs]
        bars = ax3.bar(organs, final_conc, color=colors[2:10], alpha=0.8)
        ax3.set_xlabel('Organs')
        ax3.set_ylabel('Final Concentration (mg/L)')
        ax3.set_title('Organ Distribution at 48h')
        plt.xticks(rotation=45)
        for bar, conc in zip(bars, final_conc):
            height = bar.get_height()
            ax3.text(bar.get_x() + bar.get_width()/2., height,
                    f'{conc:.2f}', ha='center', va='bottom', fontsize=8)
        
        # 4. Absorption profile
        ax4 = plt.subplot(4, 3, 4)
        abs_data = results['absorption_data']
        ax4.plot(abs_data['time_points'], abs_data['absorption_rate'], 
                color=colors[1], linewidth=2, label='Absorption Rate')
        ax4_twin = ax4.twinx()
        ax4_twin.plot(abs_data['time_points'], abs_data['absorbed_amount'], 
                     color=colors[2], linewidth=2, label='Cumulative Absorption')
        ax4.set_xlabel('Time (hours)')
        ax4.set_ylabel('Absorption Rate (mg/h)', color=colors[1])
        ax4_twin.set_ylabel('Cumulative Absorption (mg)', color=colors[2])
        ax4.set_title('Absorption Profile')
        ax4.grid(True, alpha=0.3)
        
        # 5. Organ exposure over time
        ax5 = plt.subplot(4, 3, 5)
        major_organs = ['liver', 'kidney', 'brain', 'muscle']
        for i, organ in enumerate(major_organs):
            ax5.plot(time_points, results['concentrations'][organ], 
                    color=colors[i+3], linewidth=2, label=organ.title())
        ax5.set_xlabel('Time (hours)')
        ax5.set_ylabel('Concentration (mg/L)')
        ax5.set_title('Major Organ Concentration-Time Profiles')
        ax5.legend()
        ax5.grid(True, alpha=0.3)
        
        # 6. Volume of distribution visualization
        ax6 = plt.subplot(4, 3, 6)
        organ_volumes = [self.compartments[org]['volume'] for org in major_organs]
        organ_kp = [self.compartments[org]['kp'] for org in major_organs]
        
        scatter = ax6.scatter(organ_volumes, organ_kp, 
                            c=colors[3:7], s=100, alpha=0.8)
        for i, organ in enumerate(major_organs):
            ax6.annotate(organ.title(), (organ_volumes[i], organ_kp[i]), 
                        xytext=(5, 5), textcoords='offset points', fontsize=9)
        
        ax6.set_xlabel('Organ Volume (L)')
        ax6.set_ylabel('Tissue:Plasma Partition Coefficient')
        ax6.set_title('Volume vs Distribution Coefficient')
        ax6.grid(True, alpha=0.3)
        
        # 7. Clearance pathways
        ax7 = plt.subplot(4, 3, 7)
        clearance_data = [
            self.params['pharmacokinetics']['clearance']['hepatic'],
            self.params['pharmacokinetics']['clearance']['renal'],
            0.0  # Other pathways
        ]
        labels = ['Hepatic\nCLint', 'Renal', 'Other']
        colors_cl = ['lightcoral', 'skyblue', 'lightgreen']
        
        wedges, texts, autotexts = ax7.pie(clearance_data, labels=labels, 
                                         colors=colors_cl, autopct='%1.1f%%', 
                                         startangle=90)
        ax7.set_title('Clearance Pathways Distribution')
        
        # 8. BCS classification plot
        ax8 = plt.subplot(4, 3, 8)
        drug_props = self.params['physicochemical']
        ax8.scatter(drug_props['solubility'], drug_props['permeability'] * 1e6,
                   color='red', s=200, alpha=0.8, marker='*')
        ax8.annotate('HSK7653\n(BCS III)', 
                    (drug_props['solubility'], drug_props['permeability'] * 1e6),
                    xytext=(10, 10), textcoords='offset points',
                    bbox=dict(boxstyle="round,pad=0.3", facecolor="yellow", alpha=0.7))
        
        # BCS boundaries
        ax8.axvline(x=1, color='gray', linestyle='--', alpha=0.5)
        ax8.axhline(y=1, color='gray', linestyle='--', alpha=0.5)
        ax8.text(0.1, 0.1, 'IV', transform=ax8.transAxes, fontsize=12, alpha=0.7)
        ax8.text(0.1, 0.9, 'III', transform=ax8.transAxes, fontsize=12, alpha=0.7)
        ax8.text(0.9, 0.1, 'II', transform=ax8.transAxes, fontsize=12, alpha=0.7)
        ax8.text(0.9, 0.9, 'I', transform=ax8.transAxes, fontsize=12, alpha=0.7)
        
        ax8.set_xlabel('Solubility (mg/mL)')
        ax8.set_ylabel('Permeability (×10⁻⁶ cm/s)')
        ax8.set_title('BCS Classification')
        ax8.set_xscale('log')
        ax8.set_yscale('log')
        ax8.grid(True, alpha=0.3)
        
        # 9. Summary statistics
        ax9 = plt.subplot(4, 3, 9)
        ax9.axis('off')
        
        # Calculate key metrics
        plasma_conc = results['concentrations']['plasma']
        key_metrics = self.calculate_key_metrics(results)
        
        summary_text = f"""HSK7653 PBPK Model Summary
Study: {self.params['drug_info']['study']}
Dose: {self.params['drug_info']['dose']} mg
Population: {self.params['demographics']['ethnicity']} Volunteers

Key Parameters:
• Cmax: {key_metrics['cmax']:.2f} mg/L
• Tmax: {key_metrics['tmax']:.1f} h
• AUC: {key_metrics['auc']:.1f} mg·h/L
• t1/2: {key_metrics['terminal_half_life']:.1f} h
• CL: {key_metrics['clearance']:.2f} L/h
• Vd: {key_metrics['vd']:.2f} L/kg

BCS Classification: {self.params['physicochemical']['bcs_class']}
Bioavailability: {self.params['pharmacokinetics']['bioavailability']*100:.1f}%
Systemic Clearance: {self.params['pharmacokinetics']['clearance']['systemic']} L/h
"""
        
        ax9.text(0.05, 0.95, summary_text, transform=ax9.transAxes, 
                fontsize=10, verticalalignment='top', 
                bbox=dict(boxstyle="round,pad=0.5", facecolor="lightblue", alpha=0.8))
        
        # 10. Metabolite formation
        ax10 = plt.subplot(4, 3, 10)
        # Simulate metabolite formation
        parent_fraction = np.exp(-0.693 / 8 * time_points)  # Estimated
        metabolite_fraction = 1 - parent_fraction
        ax10.plot(time_points, parent_fraction, label='Parent Drug', linewidth=2)
        ax10.plot(time_points, metabolite_fraction, label='Metabolites', linewidth=2)
        ax10.set_xlabel('Time (hours)')
        ax10.set_ylabel('Fraction')
        ax10.set_title('Drug vs Metabolite Profile')
        ax10.legend()
        ax10.grid(True, alpha=0.3)
        
        # 11. Population variability
        ax11 = plt.subplot(4, 3, 11)
        # Simulate population variability
        n_subjects = 100
        variability = np.random.normal(1, 0.2, n_subjects)
        cl_variability = self.params['pharmacokinetics']['clearance']['systemic'] * variability
        ax11.hist(cl_variability, bins=20, alpha=0.7, color='purple')
        ax11.axvline(self.params['pharmacokinetics']['clearance']['systemic'], 
                    color='red', linestyle='--', linewidth=2)
        ax11.set_xlabel('Clearance (L/h)')
        ax11.set_ylabel('Frequency')
        ax11.set_title('Population Clearance Variability')
        
        # 12. Dose proportionality
        ax12 = plt.subplot(4, 3, 12)
        doses = [1, 2.5, 5, 10, 20]  # mg
        base_auc = results['plasma_auc']
        aucs = [base_auc * (d/5) for d in doses]  # Linear PK assumption
        ax12.plot(doses, aucs, 'o-', color='green', linewidth=2, markersize=8)
        ax12.set_xlabel('Dose (mg)')
        ax12.set_ylabel('AUC (mg·h/L)')
        ax12.set_title('Dose Proportionality')
        ax12.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig(f"{save_path}/hsk7653_full_pbpk_comprehensive.png", 
                   dpi=300, bbox_inches='tight')
        plt.show()
        
        return f"{save_path}/hsk7653_full_pbpk_comprehensive.png"
    
    def run_sensitivity_analysis(self, dose: float = 5.0) -> Dict[str, Any]:
        """Run sensitivity analysis for key parameters"""
        
        base_result = self.run_full_pbpk_simulation(dose)
        
        # Parameters to test
        parameters = {
            'clearance': [0.3, 0.4, 0.5],  # L/h
            'vd': [0.4, 0.54, 0.7],       # L/kg
            'fa': [0.9, 0.998, 1.0],      # fraction
            'ka': [1.5, 2.5, 3.5]         # 1/h
        }
        
        sensitivity_results = {}
        
        for param_name, values in parameters.items():
            param_results = []
            for value in values:
                # Create modified parameters
                modified_params = self.params.copy()
                if param_name == 'clearance':
                    modified_params['pharmacokinetics']['clearance']['systemic'] = value
                elif param_name == 'vd':
                    # Adjust organ volumes proportionally
                    scale_factor = value / 0.54
                    for organ in modified_params['organ_distribution']:
                        modified_params['organ_distribution'][organ]['volume'] *= scale_factor
                elif param_name == 'fa':
                    modified_params['pharmacokinetics']['absorption']['fa'] = value
                elif param_name == 'ka':
                    modified_params['pharmacokinetics']['absorption']['ka'] = value
                
                # Run simulation with modified parameters
                original_params = self.params
                self.params = modified_params
                result = self.run_full_pbpk_simulation(dose)
                metrics = self.calculate_key_metrics(result)
                param_results.append({
                    'value': value,
                    'cmax': metrics['cmax'],
                    'auc': metrics['auc'],
                    'half_life': metrics['terminal_half_life']
                })
                self.params = original_params
            
            sensitivity_results[param_name] = param_results
        
        return sensitivity_results

def main():
    """Main execution function"""
    print("=" * 70)
    print("HSK7653 FULL PBPK MODEL SIMULATION")
    print("Study HSK7653-101-5 mg")
    print("=" * 70)
    
    # Initialize model
    model = FullPBPKHSK7653()
    
    # Run full simulation
    print("\n1. Running Full PBPK Simulation...")
    results = model.run_full_pbpk_simulation()
    
    print("\n2. Key Results:")
    print("-" * 40)
    key_metrics = model.calculate_key_metrics(results)
    print(f"Cmax: {key_metrics['cmax']:.2f} mg/L")
    print(f"Tmax: {key_metrics['tmax']:.1f} hours")
    print(f"AUC: {key_metrics['auc']:.1f} mg·h/L")
    print(f"Terminal Half-life: {key_metrics['terminal_half_life']:.1f} hours")
    print(f"Total Clearance: {key_metrics['clearance']:.2f} L/h")
    print(f"Volume of Distribution: {key_metrics['vd']:.2f} L/kg")
    
    print("\n3. Organ Exposure Summary:")
    print("-" * 40)
    for organ, exposure in results['organ_exposure'].items():
        if organ != 'gut_lumen' and exposure > 0:
            print(f"{organ.title()}: {exposure:.1f} mg·h/L")
    
    print("\n4. Generating Comprehensive Visualization...")
    chart_path = model.generate_comprehensive_plots(results)
    print(f"📊 Charts saved to: {chart_path}")
    
    print("\n5. Running Sensitivity Analysis...")
    sensitivity = model.run_sensitivity_analysis()
    
    print("\n6. Sensitivity Results:")
    print("-" * 40)
    for param, results_list in sensitivity.items():
        print(f"\n{param.upper()}:")
        for result in results_list:
            print(f"  {param}={result['value']}: AUC={result['auc']:.1f}, t1/2={result['half_life']:.1f}h")
    
    print("\n" + "=" * 70)
    print("✅ HSK7653 Full PBPK Model Simulation Complete!")
    print("=" * 70)

if __name__ == "__main__":
    main()